import argparse
import sys
import os
import gzip
import mmap
import numpy as np
import time
from multiprocessing import Manager
import multiprocessing as mp
from deepsignal3.utils.process_utils import get_motif_seqs
from deepsignal3.utils.process_utils import get_refloc_of_methysite_in_motif
from deepsignal3.utils import bam_reader

neg_num = mp.Value('i', 0)

def read_feature(input,feature_Q,batch_size=1,qsize_limit=40):
    print('begin read feature', flush=True)
    fea_list=[]
    with open(input, 'r') as f:
        for line in f:
            while feature_Q.qsize()>=qsize_limit:
                time.sleep(0.01)
            fea_list.append(line)         
            if len(fea_list)>=batch_size:
                feature_Q.put(fea_list)
                fea_list=[]
    if len(fea_list)>0:
        feature_Q.put(fea_list)

def get_q2tloc_from_cigar(r_cigar_tuple, strand, seq_len):
    """
    insertion: -1, deletion: -2, mismatch: -3
    :param r_cigar_tuple: pysam.alignmentSegment.cigartuples
    :param strand: 1/-1 for fwd/rev
    :param seq_len: read alignment length
    :return: query pos to ref pos
    """
    fill_invalid = -2
    # get each base calls genomic position
    q_to_r_poss = np.full(seq_len + 1, fill_invalid, dtype=np.int32)
    # process cigar ops in read direction
    curr_r_pos, curr_q_pos = 0, 0
    cigar_ops = r_cigar_tuple if strand == 1 else r_cigar_tuple[::-1]
    for op, op_len in cigar_ops:
        if op == 1:
            # inserted bases into ref
            for q_pos in range(curr_q_pos, curr_q_pos + op_len):
                q_to_r_poss[q_pos] = -1
            curr_q_pos += op_len
        elif op in (2, 3):
            # deleted ref bases
            curr_r_pos += op_len
        elif op in (0, 7, 8):
            # aligned bases
            for op_offset in range(op_len):
                q_to_r_poss[curr_q_pos + op_offset] = curr_r_pos + op_offset
            curr_q_pos += op_len
            curr_r_pos += op_len
        elif op == 6:
            # padding (shouldn't happen in mappy)
            pass
    q_to_r_poss[curr_q_pos] = curr_r_pos
    if q_to_r_poss[-1] == fill_invalid:
        raise ValueError(
            (
                "Invalid cigar string encountered. Reference length: {}  Cigar "
                + "implied reference length: {}"
            ).format(seq_len, curr_r_pos)
        )
    return q_to_r_poss

def process_read(feature_Q,output_Q,bam_index,batch_size=1,time_wait=0.01,output_limit=40):
    fea_list=[]
    while True:
        while feature_Q.empty():
            time.sleep(time_wait)
        read_data = feature_Q.get()
        if read_data == "kill":
            feature_Q.put("kill")
            break
        for read in read_data:
            while output_Q.qsize()>=output_limit:
                time.sleep(0.01)
            words=read.strip().split('\t')
            read_name=words[0]
            reference_name=words[1]
            if reference_name!='chr2' and reference_name!='chr4':
                continue
            ref_start=int(words[2])

            #key='||'.join(words[:2])
            try:
                for bam_read in bam_index.get_alignments(read_name):
                    if bam_read.is_supplementary or bam_read.is_secondary:
                        continue
                    if bam_read.reference_name != reference_name:
                        continue
                    if bam_read.reference_start != ref_start:
                        continue
                    qalign_start = bam_read.query_alignment_start
                    qalign_end = bam_read.query_alignment_end
                    seq_len=bam_read.infer_query_length()
                    strand_code = 0 if bam_read.is_reverse else 1
                    strand = "-" if strand_code == 0 else "+"
                    cigar_tuples = bam_read.cigartuples
                    ref_end = bam_read.reference_end
                    if bam_read.is_reverse:
                        seq_start = seq_len - qalign_end
                        seq_end = seq_len - qalign_start
                    else:
                        seq_start = qalign_start
                        seq_end = qalign_end
                    q_to_r_poss = get_q2tloc_from_cigar(
                        cigar_tuples, strand_code, (seq_end - seq_start)
                    )
                    ref_readlocs = dict()
                    ref_poss = []
                    pred_pos = []
                    pred_pos = [int(x) for x in words[5].split(",")]
                    for loc_in_read in pred_pos:
                        if seq_start <= loc_in_read < seq_end:
                            offset_idx = loc_in_read - seq_start
                            if q_to_r_poss[offset_idx] != -1:
                                if strand == "-":
                                    # pos = '.'#loc_in_read
                                    ref_pos = ref_end - 1 - \
                                        q_to_r_poss[offset_idx]
                                else:
                                    # pos = loc_in_read
                                    ref_pos = ref_start + \
                                        q_to_r_poss[offset_idx]
                                ref_readlocs[loc_in_read] = ref_pos
                                ref_poss.append(ref_pos)
                                #pred_pos.append(loc_in_read)
                            else:
                                ref_pos = -1
                                ref_readlocs[loc_in_read] = ref_pos
                                ref_poss.append(ref_pos)
                                #pred_pos.append(loc_in_read)
                        else:
                            ref_pos = -1
                            ref_readlocs[loc_in_read] = ref_pos
                            ref_poss.append(ref_pos)
                            #pred_pos.append(loc_in_read)
                    ref_pos_text = ','.join([str(ref_readlocs[x]) for x in pred_pos])
                    fea_str = '\t'.join(words+[ref_pos_text])
                    fea_list.append(fea_str)

                    if len(fea_list)>=batch_size:
                        output_Q.put(fea_list)
                        fea_list=[]
            except Exception as e:
                #traceback.print_exc()
                continue
            
            
    if len(fea_list)>0:
        output_Q.put(fea_list)

def _write_featurestr(write_fp, featurestr_q,time_wait = 3):
    print('write_process-{} starts'.format(os.getpid()))
    output=0
    # if os.path.exists(write_fp):
    #     with open(write_fp, 'a') as wf:
    #         while True:
    #             # during test, it's ok without the sleep(time_wait)
    #             if featurestr_q.empty():
    #                 time.sleep(time_wait)
    #                 continue
    #             features_str = featurestr_q.get()
    #             if features_str == "kill":
    #                 print('output line {}'.format(output))
    #                 print('write_process-{} finished'.format(os.getpid()))
    #                 break
    #             for one_features_str in features_str:
    #                 output+=1
    #                 wf.write(one_features_str + "\n")
    #             wf.flush()
    # else:
    with open(write_fp, 'w') as wf:
        while True:
            # during test, it's ok without the sleep(time_wait)
            if featurestr_q.empty():
                time.sleep(time_wait)
                continue
            features_str = featurestr_q.get()
            if features_str == "kill":
                print('output line {}'.format(output))
                print('write_process-{} finished'.format(os.getpid()))
                break
            for one_features_str in features_str:
                output+=1
                wf.write(one_features_str + "\n")
            wf.flush()

def parse_args():
    parser = argparse.ArgumentParser("")
    parser.add_argument("--input", type=str,required=True)
    parser.add_argument("--output", type=str,required=True)
    parser.add_argument("--nproc", "-p", type=int,required=True)
    parser.add_argument("--bam", type=str, required=True)
    return parser.parse_args()


def main():
    args=parse_args()
    #flag='w'
    #if os.path.exists(args.output):
    #    flag='a'

    manager = mp.Manager()
    feature_Q = manager.Queue()
    output_Q = manager.Queue()
    bam_index = bam_reader.ReadIndexedBam(args.bam)
  
    pd=mp.Process(target=read_feature,args=(args.input,feature_Q),name="fea_reader")
    pd.daemon = True
    pd.start()
    ex_dp=args.nproc - 2
    ex_procs = []
    for i in range(ex_dp):
        pb = mp.Process(target=process_read, args=(feature_Q,output_Q,bam_index),
                          name="pb_reader")
            
        pb.daemon = True
        pb.start()
        ex_procs.append(pb)
    p_w = mp.Process(target=_write_featurestr, args=(args.output, output_Q),
                     name="writer")
    p_w.daemon = True
    p_w.start()
    pd.join()
    feature_Q.put("kill")
    for pb in ex_procs:
        pb.join()
    output_Q.put("kill")
    p_w.join()
    print('add neg number is: {}'.format(neg_num.value),flush=True)
    
        


if __name__ == '__main__':
    sys.exit(main())

